import sys
import os
import click
from datetime import datetime, timedelta
import random

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from app import create_app, db
from app.models import User, Sample, MarkerAnalysis

app = create_app()

@click.group()
def cli():
    pass

@cli.command()
def init():
    """Initialize the database."""
    with app.app_context():
        db.create_all()
        click.echo('Initialized the database.')

@cli.command()
def drop():
    """Drop all tables."""
    if click.confirm('Are you sure you want to drop all tables?'):
        with app.app_context():
            db.drop_all()
            click.echo('Dropped all tables.')

@cli.command()
@click.option('--username', prompt=True)
@click.option('--email', prompt=True)
@click.option('--password', prompt=True, hide_input=True, confirmation_prompt=True)
@click.option('--role', prompt=True, default='researcher')
def create_user(username, email, password, role):
    """Create a new user."""
    with app.app_context():
        # 检查用户名是否已存在
        if User.query.filter_by(username=username).first():
            click.echo(f'Error: Username {username} already exists.')
            return
        
        # 检查邮箱是否已存在
        if User.query.filter_by(email=email).first():
            click.echo(f'Error: Email {email} already exists.')
            return
        
        try:
            user = User(username=username, email=email, role=role)
            user.set_password(password)
            db.session.add(user)
            db.session.commit()
            click.echo(f'Successfully created user {username}')
        except Exception as e:
            db.session.rollback()
            click.echo(f'Error creating user: {str(e)}')

@cli.command()
def list_users():
    """List all users."""
    with app.app_context():
        users = User.query.all()
        if not users:
            click.echo('No users found.')
            return
        
        click.echo('\nUsers:')
        for user in users:
            click.echo(f'- {user.username} ({user.email}) [Role: {user.role}]')

@cli.command()
@click.argument('username')
def delete_user(username):
    """Delete a user by username."""
    if click.confirm(f'Are you sure you want to delete user {username}?'):
        with app.app_context():
            user = User.query.filter_by(username=username).first()
            if not user:
                click.echo(f'User {username} not found.')
                return
            
            db.session.delete(user)
            db.session.commit()
            click.echo(f'Deleted user {username}')

@cli.command()
def create_test_data():
    """Create test data for samples."""
    with app.app_context():
        # 确保有管理员用户
        admin = User.query.filter_by(username='admin').first()
        if not admin:
            click.echo('Error: Please create admin user first')
            return

        # 测试数据
        varieties = [
            '新麦26', '新麦19', '郑麦366', '周麦32', '百农207',
            '郑麦379', '周麦22', '新麦18', '郑麦7698', '周麦36'
        ]
        
        locations = [
            '新疆昌吉', '新疆伊犁', '新疆巴州', '新疆阿克苏', '新疆喀什',
            '新疆吐鲁番', '新疆哈密', '新疆克拉玛依', '新疆博州', '新疆塔城'
        ]

        # 生成50个测试样品
        for i in range(50):
            # 随机生成采集日期（最近一年内）
            collection_date = datetime.now() - timedelta(days=random.randint(0, 365))
            
            # 随机生成样品数据
            sample = Sample(
                sample_id=f'XJ{datetime.now().year}{i+1:03d}',
                variety_name=random.choice(varieties),
                collection_date=collection_date.date(),
                location=random.choice(locations),
                hardness=round(random.uniform(20, 80), 2),
                protein_content=round(random.uniform(10, 16), 2),
                moisture_content=round(random.uniform(8, 14), 2),
                notes=f'测试样品 #{i+1}',
                user_id=admin.id
            )
            
            db.session.add(sample)
        
        try:
            db.session.commit()
            click.echo('Successfully created 50 test samples')
        except Exception as e:
            db.session.rollback()
            click.echo(f'Error creating test data: {str(e)}')

@cli.command()
def clear_test_data():
    """Clear all test data from samples table."""
    if click.confirm('Are you sure you want to delete all sample data?'):
        with app.app_context():
            try:
                Sample.query.delete()
                db.session.commit()
                click.echo('Successfully cleared all sample data')
            except Exception as e:
                db.session.rollback()
                click.echo(f'Error clearing test data: {str(e)}')

@cli.command()
def create_analysis_test_data():
    """Create test data for marker analyses."""
    with app.app_context():
        # 确保有管理员用户
        admin = User.query.filter_by(username='admin').first()
        if not admin:
            click.echo('Error: Please create admin user first')
            return
        
        # 获取所有样品
        samples = Sample.query.filter_by(user_id=admin.id).all()
        if not samples:
            click.echo('Error: Please create sample test data first')
            return

        # 基因型数据
        gene_types = ['wild', 'mutation', 'null']
        
        # 示例序列数据
        sequence_templates = [
            "ATGGCAACCGTGTTCAACAATCTCCTCCTCGCGCTCCTCCTCTGCCTGCTGCCGCTCGCCACCGCCATGGCCGGTGGCGATGGC",
            "ATGGCAACCGTGTTCAACAATCTCCTCCTCGCGCTCCTCCTCTGCCTGCTGCCGCTCGCCACCGCCATGGCCGGTGGCGATGGC",
            "ATGGCAACCGTGTTCAACAATCTCCTCCTCGCGCTCCTCCTCTGCCTGCTGCCGCTCGCCACCGCCATGGCCGGTGGCGATGGC"
        ]
        
        # 备注模板
        notes_templates = [
            "序列分析显示为典型的野生型",
            "发现基因突变位点，可能影响蛋白质功能",
            "检测到基因缺失，完全丧失功能",
            "序列质量良好，结果可靠",
            "需要进一步验证结果"
        ]

        # 为每个样品生成1-3条分析记录
        for sample in samples:
            for _ in range(random.randint(1, 3)):
                # 随机生成分析日期（最近3个月内）
                analysis_date = datetime.now() - timedelta(
                    days=random.randint(0, 90),
                    hours=random.randint(0, 23),
                    minutes=random.randint(0, 59)
                )
                
                # 随机选择基因型
                pina_type = random.choice(gene_types)
                pinb_type = random.choice(gene_types)
                
                # 生成序列数据（随机选择模板并添加随机变异）
                sequence_data = random.choice(sequence_templates)
                if random.random() < 0.3:  # 30%概率添加变异
                    pos = random.randint(0, len(sequence_data)-1)
                    bases = list('ATCG')
                    bases.remove(sequence_data[pos])
                    sequence_data = sequence_data[:pos] + random.choice(bases) + sequence_data[pos+1:]
                
                # 创建分析记录
                analysis = MarkerAnalysis(
                    sample_id=sample.id,
                    pina_type=pina_type,
                    pinb_type=pinb_type,
                    grain_hardness=determine_grain_hardness(pina_type, pinb_type),
                    sequence_data=sequence_data,
                    analysis_date=analysis_date,
                    notes=random.choice(notes_templates),
                    user_id=admin.id
                )
                
                db.session.add(analysis)
        
        try:
            db.session.commit()
            click.echo('Successfully created test marker analyses')
        except Exception as e:
            db.session.rollback()
            click.echo(f'Error creating test data: {str(e)}')

def determine_grain_hardness(pina_type, pinb_type):
    """根据Pina和Pinb基因类型判断籽粒硬度"""
    if pina_type == 'wild' and pinb_type == 'wild':
        return 'soft'  # 软质
    elif pina_type == 'null' or pinb_type == 'null':
        return 'hard'  # 硬质
    elif pina_type == 'mutation' or pinb_type == 'mutation':
        return 'medium'  # 中等
    else:
        return 'unknown'  # 未知

if __name__ == '__main__':
    cli() 